from tqdm import tqdm
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
import matplotlib.pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

batch_size = 64
learning_rate = 0.001
num_epochs = 30
num_classes = 3

data_transform = transforms.Compose([
    transforms.Grayscale(),  # 保证是灰度图
    transforms.Resize((48, 48)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

train_dataset = datasets.ImageFolder(root='../data_happy_angry/train', transform=data_transform)
val_dataset = datasets.ImageFolder(root='../data_happy_angry/test', transform=data_transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)


class EmotionCNN(nn.Module):
    def __init__(self, num_classes):
        super(EmotionCNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),  # input: 1x48x48
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2),  # 32x24x24

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(2),  # 64x12x12

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(128),
            nn.MaxPool2d(2),  # 128x6x6
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 6 * 6, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x


model = EmotionCNN(num_classes).to(device)
model.load_state_dict(torch.load("emotion_model_happy_angry_MyPhoto.pth"))

from PIL import Image

def predict_image(image_path):
    model.eval()
    image = Image.open(image_path).convert('L')  # 转为灰度图
    image = data_transform(image).unsqueeze(0).to(device)
    output = model(image)
    _, predicted = torch.max(output.data, 1)
    return train_dataset.classes[predicted.item()]

def predict_images_in_folder(folder_path):
    # 遍历文件夹下的所有文件
    results = {}
    for filename in os.listdir(folder_path):
        if filename.endswith('.jpg') or filename.endswith('.png'):  # 只处理jpg和png文件
            image_path = os.path.join(folder_path, filename)
            prediction = predict_image(image_path)
            results[filename] = prediction
    return results

# 示例
folder_path = r'../valdata/happy_faces'  # 替换为你自己的文件夹路径
predictions = predict_images_in_folder(folder_path)


count =0
total = 0
# 输出结果
for filename, prediction in predictions.items():
    total +=1
    if prediction =='happy':
        count+=1
    print(f"{filename}: {prediction}")
print(total)
print(count)


# 示例
folder_path = r'../valdata/angry_faces'   # 替换为你自己的文件夹路径
predictions = predict_images_in_folder(folder_path)

count_angry =0
total_angry = 0
# 输出结果
for filename, prediction in predictions.items():
    total_angry +=1
    if prediction =='angry':
        count_angry+=1
    print(f"{filename}: {prediction}")
print(total_angry)
print(count_angry)
print(count_angry/total_angry*100)

# 示例
folder_path = r'../valdata/other_faces'   # 替换为你自己的文件夹路径
predictions = predict_images_in_folder(folder_path)

count_other =0
total_other = 0
# 输出结果
for filename, prediction in predictions.items():
    total_other +=1
    if prediction =='other':
        count_other+=1
    print(f"{filename}: {prediction}")
print(total_other)
print(count_other)
print(count_other/total_other*100)
# 这个代码主要检验文件夹下照片情绪是否判断正确

# 先用这个代码检验原始数据集训练出的模型在我的照片中的效果（初始模型对我的表情的识别 ）  G:\myphoto\angry_faces  G:\myphoto\happy_faces G:\myphoto\other_faces
# 发现 121 张happy照片中，只有49张识别为happy
#  61 张angry照片中，47张识别为angry
#  88 张other照片中，77张识别为other


# 再用这个代码检验 原始数据集+我的数据集 训练出的模型在我的 新 照片中的效果（强化后的模型对我表情的识别 ）  G:\myphoto_new\angry_faces  G:\myphoto_new\happy_faces
# 发现 121 张happy照片中，只有49张识别为happy
#  67 张angry照片中，65张识别为angry
#  86 张other照片中，83张识别为other